package nl.rug.jvleusen.datastructurelib.tree;

import nl.rug.jvleusen.datastructurelib.ITree;

/**
 * The AVLTree implementation is an adaption of the {@link BinaryTree}
 * implementation, the difference is that if the tree is imbalanced by an insert
 * or remove action, it is restructured to ensure a binarySeach will always have
 * <tt> O(log(n))</tt> complexity.
 *
 * @author Jeroen van Leusen
 * @see BinaryTree
 * @see ITree
 */
public class AVLTree<E extends Comparable<E>> implements ITree<E, AVLTree.AVLTreeNode<E>> {

    public AVLTreeNode<E> root = null;
    private int size = 0;

    /**
     * {@inheritDoc}
     */
    public boolean isEmpty() {
        return root == null;
    }

    /**
     * {@inheritDoc}
     */
    public int height() {
        if (root == null) {
            return -1;
        } else {
            return root.height();
        }
    }

    /**
     * {@inheritDoc}
     */
    public AVLTreeNode<E> getRoot() {
        return root;
    }

    /**
     * @return the number of elements in this tree.
     */
    public int size() {
        return size;
    }

    /**
     * Checks whether there is a node containing <tt>obj</tt> within this tree.
     *
     * @param obj The object to look for
     * @return Whether obj could be found.
     */
    public boolean contains(E obj) {
        return binarySearch(obj) != null;
    }

    /**
     * Preforms a binary search on this tree, returning the node if it finds a
     * node which {@code getData().compareTo(data) == 0} is true for. The
     * complexity is {@code O(log(n))}, which is enforced by the rebalancing
     * during inserts/removals.
     *
     * @param data The data to match to.
     * @return The node that contains the data which matches with the given
     * data, if it exists, returns null otherwise.
     * @throws IllegalArgumentException If you try to look for an element
     * containing a null element, because they cannot be compared.
     */
    public AVLTreeNode<E> binarySearch(E data) {
        if (data == null) {
            throw new IllegalArgumentException("null is not a valid element.");
        }

        if (root == null) {
            return null;
        }

        AVLTreeNode<E> node = root;
        int compare;
        while ((compare = node.getData().compareTo(data)) != 0) {
            if (compare > 0) {
                if (node.left() != null) {
                    node = node.left();
                } else {
                    return null;
                }
            } else {
                if (node.right() != null) {
                    node = node.right();
                } else {
                    return null;
                }
            }
        }

        return node;
    }

    /**
     * Attempts to remove the node with the given data, if it is found the tree
     * is rebalanced.
     *
     * @param data element to look for.
     * @return Whether something was successfully removed.
     */
    public boolean remove(E data) {
        if (root == null) {
            return false;
        } else {
            if (root.getData().compareTo(data) == 0) {
                if (root.left() != null && root.right() == null) {
                    //Just left, so root can be replaced by left.
                    root = root.left();
                    root.setParent(null);
                } else if (root.left() == null && root.right() != null) {
                    //Ditto for the right.
                    root = root.right();
                    root.setParent(null);
                } else if (root.left() == null && root.right() == null) {
                    root = null;
                } else {
                    AVLTreeNode<E> min = root.right().minValue();
                    root.data = min.getData();
                    AVLTreeNode<E> remove = root.right().remove(root.data, this);
                    balance(remove);
                }

                size--;
                return true;
            } else {
                AVLTreeNode<E> parentRemoved = root.remove(data, this);
                if (parentRemoved != null) {
                    balance(parentRemoved);
                    size--;
                }
                return parentRemoved != null;
            }
        }
    }

    /**
     * Attempts to insert a node with the given data according to the rules of
     * the binary search tree. If the node doesn't yet exist then a new node is
     * added and the tree is rebalanced.
     *
     * @param data element to insert.
     * @return Whether the element was successfully added.
     * @throws IllegalArgumentException If you try to insert null.
     */
    public boolean insert(E data) {
        if (data == null) {
            throw new IllegalArgumentException("This implementation does not take null objects.");
        } else if (root == null) {
            root = new AVLTreeNode<E>(data, null);
            size++;
            return true;
        } else {
            if (insert(data, root)) {
                size++;
                return true;
            } else {
                return false;
            }
        }
    }

    private boolean insert(E data, AVLTreeNode<E> node) {
        int compare = node.getData().compareTo(data);
        if (compare == 0) {
            //Stuff already in the tree, return false
            return false;
        } else {
            if (compare > 0) {
                AVLTreeNode<E> left = node.left();
                if (left == null) {
                    node.setLeft(new AVLTreeNode<E>(data, node));
                } else {
                    return insert(data, left);
                }
            } else {
                //Check right
                AVLTreeNode<E> right = node.right();
                if (right == null) {
                    node.setRight(new AVLTreeNode<E>(data, node));
                } else {
                    return insert(data, right);
                }
            }
            //Rebalance the tree, since we added something new to this node.
            balance(node);
            return true;
        }
    }

    /*
     * Internal balancing methods.
     */
    /**
     * Iterates upwards in the tree, balancing the tree if a specific subtree is
     * imbalanced. It should be called on the lowest node affected by the latest
     * change.
     *
     * @param node
     */
    private void balance(AVLTreeNode<E> node) {
        if (!node.isLeaf()) {
            calcHeight(node);
        }
        while (node != root) { //While the node isn't the exact same as the root of the tree...
            node = node.parent();
            calcHeight(node);
            if (!isBalanced(node)) {
                //preform a trinode restructuring at node's tallest grandchild, for it to be imbalanced there has to be one.
                AVLTreeNode<E> node2 = tallerChild(tallerChild(node));
                node = restructure(node2); //tri-node restructure from the parent.
                calcHeight(node.left());
                calcHeight(node.right());
                calcHeight(node);
            }
        }
    }

    private AVLTreeNode<E> tallerChild(AVLTreeNode<E> node) {
        if (height(node.left()) > height(node.right())) { //left is taller.
            return node.left();
        } else if (height(node.left()) < height(node.right())) {
            return node.right();
        } else { //Children have equal height.
            if (node == root) { //We cannot use the next trick if its the root.
                return node.left();
            }

            if (node == node.parent().left()) {
                return node.left();
            } else {
                return node.right();
            }
        }
    }

    private void calcHeight(AVLTreeNode<E> node) {
        node.height = 1 + Math.max(height(node.left()), height(node.right()));
    }

    private int height(AVLTreeNode<E> node) {
        if (node == null) {
            return -1;
        } else {
            return node.height();
        }
    }

    private boolean isBalanced(AVLTreeNode<E> node) {
        int bf = height(node.left()) - height(node.right());
        return ((-1 <= bf) && (bf <= 1)); //the balance is not below -1 or above 1
    }

    /**
     * Performs a tri-node restructuring. Assumes the nodes are in one of
     * following configurations:
     *
     * <pre>
     *          z=c       z=c        z=a         z=a
     *         /  \      /  \       /  \        /  \
     *       y=b  t4   y=a  t4    t1  y=c     t1  y=b
     *      /  \      /  \           /  \         /  \
     *    x=a  t3    t1 x=b        x=b  t4       t2 x=c
     *   /  \          /  \       /  \             /  \
     *  t1  t2        t2  t3     t2  t3           t3  t4
     * </pre>
     *
     * @return the new root of the restructured subtree
     */
    private AVLTreeNode<E> restructure(AVLTreeNode<E> x) {
        AVLTreeNode<E> a, b, c, t1, t2, t3, t4;
        AVLTreeNode<E> y = x.parent();
        AVLTreeNode<E> z = y.parent();
        boolean xLeft = (x == y.left());
        boolean yLeft = (y == z.left());
        if (xLeft && yLeft) {
            a = x;
            b = y;
            c = z;
            t1 = a.left();
            t2 = a.right();
            t3 = b.right();
            t4 = c.right();
        } else if (!xLeft && yLeft) {
            a = y;
            b = x;
            c = z;
            t1 = a.left();
            t2 = b.left();
            t3 = b.right();
            t4 = c.right();
        } else if (xLeft && !yLeft) {
            a = z;
            b = x;
            c = y;
            t1 = a.left();
            t2 = b.left();
            t3 = b.right();
            t4 = c.right();
        } else { // right-right 
            a = z;
            b = y;
            c = x;
            t1 = a.left();
            t2 = b.left();
            t3 = c.left();
            t4 = c.right();
        }
        if (z == root) {
            b.setParent(root.parent());
            root = b;
        } else {
            AVLTreeNode<E> zParent = z.parent();
            if (z == zParent.left()) {
                b.setParent(zParent);
                zParent.setLeft(b);
            } else { // z was the right child.
                b.setParent(zParent);
                zParent.setRight(b);
            }
        }
        b.setLeft(a);
        a.setParent(b);
        b.setRight(c);
        c.setParent(b);
        a.setLeft(t1);
        if (t1 != null) {
            t1.setParent(a);
        }
        a.setRight(t2);
        if (t2 != null) {
            t2.setParent(a);
        }
        c.setLeft(t3);
        if (t3 != null) {
            t3.setParent(c);
        }
        c.setRight(t4);
        if (t4 != null) {
            t4.setParent(c);
        }

        return b;
    }

    protected AVLTreeNode<E> removeNodeFromParent(AVLTreeNode<E> node) {
        AVLTreeNode<E> parent = node.parent();
        if (node == parent.left()) {
            AVLTreeNode<E> newNode = (node.left() != null) ? node.left() : node.right();
            parent.left = newNode;
            if (newNode != null) {
                newNode.setParent(parent);
            }
        } else {
            AVLTreeNode<E> newNode = (node.left() != null) ? node.left() : node.right();
            parent.right = newNode;
            if (newNode != null) {
                newNode.setParent(parent);
            }
        }
        return parent;
    }

    /**
     * Implementation of {@link TreeNode} that provides the required methods for {@link AVLTree}
     */
    public static class AVLTreeNode<E extends Comparable<E>> implements ITree.TreeNode<E> {

        private AVLTreeNode<E> parent;
        private AVLTreeNode<E> left;
        private AVLTreeNode<E> right;
        private E data;
        private int height;

        /**
         * Creates a node with the given data and parent initialized, with left
         * and right being null.
         *
         * @param data data to initialize with
         * @param parent parent node to initialize with.
         */
        public AVLTreeNode(E data, AVLTreeNode<E> parent) {
            this(data, parent, null, null);
        }

        /**
         * Creates a node with the given initial values.
         *
         * @param data element to initialize with.
         * @param parent parent node to initialize with.
         * @param left node to set as left child.
         * @param right node to set as right child.
         */
        public AVLTreeNode(E data, AVLTreeNode<E> parent, AVLTreeNode<E> left, AVLTreeNode<E> right) {
            this.data = data;
            this.parent = parent;
            this.left = left;
            this.right = right;

            height = 0;

            if (left != null) {
                height = Math.max(height, 1 + left.height());
            }

            if (right != null) {
                height = Math.max(height, 1 + right.height());
            }
        }

        /**
         * {@inheritDoc}
         */
        public E getData() {
            return data;
        }

        /**
         * @return The left child.
         */
        public AVLTreeNode<E> left() {
            return left;
        }

        protected void setLeft(AVLTreeNode<E> left) {
            this.left = left;
        }

        /**
         * @return The right child.
         */
        public AVLTreeNode<E> right() {
            return right;
        }

        protected void setRight(AVLTreeNode<E> right) {
            this.right = right;
        }

        /**
         * @return The parent node of this node.
         */
        public AVLTreeNode<E> parent() {
            return parent;
        }

        protected void setParent(AVLTreeNode<E> parent) {
            this.parent = parent;
        }

        /**
         * {@inheritDoc}
         */
        public boolean isLeaf() {
            return left == null && right == null;
        }

        /**
         * {@inheritDoc}
         */
        public int height() {
            return height;
        }

        protected AVLTreeNode<E> remove(E value, AVLTree<E> tree) {
            int compare = getData().compareTo(value);
            if (compare > 0) {
                if (left != null) {
                    return left.remove(value, tree);
                } else {
                    return null;
                }
            } else if (compare < 0) {
                if (right != null) {
                    return right.remove(value, tree);
                } else {
                    return null;
                }
            } else {
                AVLTreeNode<E> removedParent = null;
                if (left != null && right != null) {
                    AVLTreeNode<E> minNode = right.minValue();
                    this.data = minNode.getData();
                    removedParent = tree.removeNodeFromParent(minNode);
                } else {
                    removedParent = tree.removeNodeFromParent(this);
                }
                return removedParent;
            }
        }

        private AVLTreeNode<E> minValue() {
            if (left == null) {
                return this;
            } else {
                return left.minValue();
            }
        }

        @Override
        public String toString() {
            return String.format("AVLTreeNode{value=%s,height=%d}", data, height);
        }
    }
}
